import os
# import random
import sys
import threading
# from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed


from PIL import Image
from PIL import ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True

Image.MAX_IMAGE_PIXELS = None  # 禁用解压缩炸弹限制



if not os.path.exists("mozcjpeg.exe"):
    print("未找到mozcjpeg")
    input("按任意键结束")
    sys.exit()
else:
    print("已找到mozcjpeg")

if not os.path.exists("gifsicle.exe"):
    print("未找到gifsicle")
    input("按任意键结束")
    sys.exit()
else:
    print("已找到gifsicle")

def pd_img(infile):  # 判断图片格式，imghdr不准确
    try:
        im = Image.open(infile, "r")
    except BaseException:
        return None
    return "."+im.format


ysquality1 = 65  # jpegtran质量，实际上是PIL转换质量，jpegtran是无损
ysquality2 = 70  # mozjpeg质量
smooth = 25  # 平滑度

min_ys = 200  # 小于该值图片不压缩，单位kb
maxjd_ys = 1000  # 大于该值图片压缩狠一点，单位kb
maxjd_ys2 = 2000  # 大于该值图片压缩再狠一点，单位kb
after_ys = "ysy"

max_Resolution = 1921  # 分辨率太大的话把图片缩小
zoom_ratio = 0.7  # 缩放比例

tpgs = ['.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG', '.bmp', '.BMP', '.WEBP', '.webp']
gif_tpgs = ['.gif', '.GIF']
mutex = threading.Lock()  # 文件互斥锁


def pilys2(infile, code=0, ysquality=ysquality2,  ig_code=1):  # jpg模式2压缩
    infile=winapi_path(infile)
    if not os.path.exists(infile):  # 判断图片还在不在
        print("图片消失了")
        return 0

    if "pl2temp" in os.path.splitext(infile)[0]:  # 如果是临时图片不压缩还要删除
        try:
            os.remove(infile)
        except BaseException:
            print("临时文件无法删除")
            return 1
        return 0

    if os.path.getsize(infile) / 1024 < min_ys:  # 太小了不压缩
        return 0

    newfile_name1 = ""
    newfile_name2 = os.path.splitext(infile)[0] + "_pl2temp" + ".jpeg"  # 临时文件
    newfile_name3 = os.path.splitext(infile)[0] + "_pl2temp2" + ".jpeg"  # 临时文件2
    if code == 1:
        newfile_name1 = os.path.splitext(infile)[0] + ".jpeg"
    else:
        newfile_name1 = os.path.splitext(infile)[0] + "_pl2" + ".jpeg"
    pdimgvalue = pd_img(infile)
    if pdimgvalue is not None and pdimgvalue in tpgs and pdimgvalue != '':
        pass
    else:
        if ig_code == 0:  # 忽视代码，0为忽视，1为不忽视
            return 0
        print(infile + "不是支持的图片")
        return 1  # 不是图片不压缩

    if os.path.exists(newfile_name2):
        try:
            os.remove(newfile_name2)
        except BaseException:
            print("临时文件无法删除")
            return 1
    if os.path.exists(newfile_name3):
        try:
            os.remove(newfile_name3)
        except BaseException:
            print("临时文件无法删除")
            return 1

    im = Image.open(infile)
    im = im.convert('RGB')

    width, height = im.size  # 图片的长和宽
    if width >= max_Resolution or height >= max_Resolution:  # 图片太大了的话缩小尺寸
        im = im.resize(((int)(width * zoom_ratio), (int)(height * zoom_ratio)), Image.LANCZOS)
    im.save(newfile_name3, "JPEG", quality=98)

    if os.path.getsize(newfile_name3) / 1024 < min_ys:  # 太小了不压缩
        if code == 1:
            if os.path.exists(infile):  # 如果文件存在
                try:
                    os.remove(infile)  # 删除原文件
                    if os.path.exists(newfile_name1):
                        os.remove(newfile_name1)
                    if not os.path.exists(infile) and os.path.exists(newfile_name3):
                        os.rename(newfile_name3, newfile_name1)  # 重命名新文件为原文件名
                        write_scinform(newfile_name1, after_ys)
                except BaseException:
                    print("删除原文件错误")
                    return 1
            else:
                if os.path.exists(newfile_name1):
                    os.remove(newfile_name1)
                if not os.path.exists(infile) and os.path.exists(newfile_name3):
                    os.rename(newfile_name3, newfile_name1)  # 重命名新文件为原文件名
                    write_scinform(newfile_name1, after_ys)
        else:
            if os.path.exists(newfile_name3):
                os.rename(newfile_name3, newfile_name1)  # 重命名新文件为原文件名
                write_scinform(newfile_name1, after_ys)
        if os.path.exists(newfile_name2):
            try:
                os.remove(newfile_name2)
            except BaseException:
                print("临时文件无法删除")
        return 0
    if os.path.getsize(newfile_name3) / 1024 > maxjd_ys2:  # 图片太大的话降低质量
        ysquality = int(ysquality / 1.7)
    elif os.path.getsize(newfile_name3) / 1024 > maxjd_ys:
        ysquality = int(ysquality / 1.4)

    cmd = "mozcjpeg -tune-ssim -quality " + str(ysquality) + " -smooth " + str(smooth) + " " + "< \"" + newfile_name3 + "\" >" + " " + "\"" + newfile_name2 + "\""
    if os.system(cmd) != 0:
        print(infile + "mozcjpeg压缩失败")
        if os.path.exists(newfile_name2):
            os.remove(newfile_name2)
        if os.path.exists(newfile_name3):
            os.remove(newfile_name3)
        if ig_code == 0:  # 忽视代码，0为忽视，1为不忽视
            return 0
        write_erro(infile)
        return 1

    if os.path.exists(newfile_name3):
        os.remove(newfile_name3)

    if code == 2 and os.path.exists(newfile_name2):
        os.remove(newfile_name2)
    elif code == 1:
        if os.path.exists(infile):  # 如果文件存在
            try:
                os.remove(infile)  # 删除原文件
                if os.path.exists(newfile_name1) and os.path.exists(newfile_name2):
                    os.remove(newfile_name1)
                if not os.path.exists(infile) and os.path.exists(newfile_name2):
                    os.rename(newfile_name2, newfile_name1)  # 重命名新文件为原文件名
            except BaseException:
                if os.path.exists(newfile_name2):
                    try:
                        os.remove(newfile_name2)
                    except BaseException:
                        print("临时文件无法删除")
                if os.path.exists(newfile_name3):
                    try:
                        os.remove(newfile_name3)
                    except BaseException:
                        print("临时文件无法删除")
                print("删除原文件错误")
                return 1
    write_scinform(newfile_name1, after_ys)
    return 0


def gif_img_ps(infile, code=0,ig_code=1):  # jpg模式2压缩
    
    #infile=winapi_path(infile)
    if not os.path.exists(infile):  # 判断图片还在不在
        print("图片消失了")
        return 0

    if "pl2temp" in os.path.splitext(infile)[0]:  # 如果是临时图片不压缩还要删除
        try:
            os.remove(infile)
        except BaseException:
            print("临时文件无法删除")
            return 1
        return 0

    if os.path.getsize(infile) / 1024 < min_ys:  # 太小了不压缩
        return 0

    newfile_name1 = ""
    newfile_name2 = os.path.splitext(infile)[0] + "_pl2temp" + ".gif"  # 临时文件
    
    if code == 1:
        newfile_name1 = os.path.splitext(infile)[0] + ".gif"
    else:
        newfile_name1 = os.path.splitext(infile)[0] + "_pl2" + ".gif"
    pdimgvalue = pd_img(infile)
    if pdimgvalue is not None and pdimgvalue in gif_tpgs and pdimgvalue != '':
        pass
    else:
        if ig_code == 0:  # 忽视代码，0为忽视，1为不忽视
            return 0
        print(infile + "不是支持的图片")
        return 1  # 不是图片不压缩

    if os.path.exists(newfile_name2):
        try:
            os.remove(newfile_name2)
        except BaseException:
            print("临时文件无法删除")
            return 1

    cmd = "gifsicle " + " --lossy=70 --color-method blend-diversity  -O3  "  +" \""+infile+ "\" "+" > "+" \""+ newfile_name2+ "\" "
    if os.system(cmd) != 0:
        print(infile + " gif压缩失败")
        if os.path.exists(newfile_name2):
            os.remove(newfile_name2)
        if ig_code == 0:  # 忽视代码，0为忽视，1为不忽视
            return 0
        write_erro(infile)
        return 1

    if code == 2 and os.path.exists(newfile_name2):
        os.rename(newfile_name2, newfile_name1)  # 重命名新文件为原文件名
    elif code == 1:
        if os.path.exists(infile):  # 如果文件存在
            try:
                os.remove(infile)  # 删除原文件
                if not os.path.exists(infile) and os.path.exists(newfile_name2):
                    os.rename(newfile_name2, newfile_name1)  # 重命名新文件为原文件名
            except BaseException:
                if os.path.exists(newfile_name2):
                    try:
                        os.remove(newfile_name2)
                    except BaseException:
                        print("临时文件无法删除")
                print("删除原文件错误")
                return 1
    write_scinform(newfile_name1, after_ys)
    return 0


def write_erro(mydir):  # 写入错误文件名
    if not os.path.exists(os.path.abspath(mydir)):
        return 1
    errodir = os.path.split(os.path.abspath(mydir))[0] + "/erro"
    # 若文件不存在，系统自动创建。'a'表示可连续写入到文件，保留原内容
    mutex.acquire()

    f = open(errodir, 'a', encoding='utf-8')
    if mydir.startswith(u"\\\\?\\"):
        mydir=mydir.replace(u"\\\\?\\","")
    f.write("\n")  # 换行
    f.write(mydir)
    f.close()
    mutex.release()
    return 0


def write_scinform(infile, flag):
    if not os.path.exists(os.path.abspath(infile)):
        return 1
    mutex.acquire()
    f = open(os.path.split(os.path.abspath(infile))[0] + "\\" + flag, 'a', encoding='utf-8')
    if infile.startswith(u"\\\\?\\"):
        infile=infile.replace(u"\\\\?\\","")
    f.write("\n")  # 换行
    f.write(infile)
    f.close()
    mutex.release()
    return 0


def write_sc(mydir, flag='123'):  # 写入成功标志
    if not os.path.exists(mydir):
        return 1
    f = open(mydir + "\\" + flag, 'a')  # 若文件不存在，系统自动创建。'a'表示可连续写入到文件，保留原内容
    f.close()
    return 0


def find_img(mydir, flag):  # 寻找该文件夹内部图片的详细地址（不含子文件夹）避开已压缩图片
    file_list1 = []
    file_list2 = []
    for folderName, subfolders, filenames in os.walk(mydir):
        for filename in filenames:
            if os.path.splitext(filename)[1] is not None and os.path.splitext(filename)[1] in tpgs and os.path.splitext(filename)[1] != '':
                if "pl2temp" not in filename:
                    file_list1.append(os.path.join(folderName, filename))
    if os.path.exists(mydir + '\\' + flag):  # 避开已压缩图片
        with open(mydir + "/" + flag, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip('\n') != '':
                    file_list2.append(line.strip('\n'))
        return list(set(file_list1) ^ set(file_list2))
    return file_list1


def find_gif_img(mydir, flag):  # 寻找该文件夹内部图片的详细地址（不含子文件夹）避开已压缩图片
    file_list1 = []
    file_list2 = []
    for folderName, subfolders, filenames in os.walk(mydir):
        for filename in filenames:
            if os.path.splitext(filename)[1] is not None and os.path.splitext(filename)[1] in gif_tpgs and os.path.splitext(filename)[1] != '':
                if "pl2temp" not in filename:
                    file_list1.append(os.path.join(folderName, filename))
    if os.path.exists(mydir + '\\' + flag):  # 避开已压缩图片
        with open(mydir + "/" + flag, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip('\n') != '':
                    file_list2.append(line.strip('\n'))
        return list(set(file_list1) ^ set(file_list2))
    return file_list1


def find_flag(mydir, flag='123'):  # 寻找标志
    for folderName, subfolders, filenames in os.walk(mydir):
        for filename in filenames:
            if filename == flag:
                return 1
    return 0



# print("0 智能判断压缩方式  1 采用png压缩（若图片非png则调用jpg压缩）")
# print("2 转换为WEBP       3 采用jpg模式1压缩（jpegtran）")
# print("4 采用jpg模式2压缩（mozcjpeg）")
def dirys(mydir, model=0, code=0, works=5, flag='123'):  # 压缩一个文件夹的图片
    img_list1 = []
    gif_img_list = []
    data = 0
    print("将采用JPEG（mozcjpeg模式）压缩" + str(mydir) + "下文件")
    if not os.path.isdir(mydir):
        return 1
    if find_flag(mydir, flag) == 1:
        print(mydir + "内部图片已压缩")
        return 1
    img_list1 = find_img(mydir, after_ys)
    gif_img_list = find_gif_img(mydir, after_ys)
    if len(img_list1) == 0 and len(gif_img_list) ==0:
        print(mydir + "没有寻找到需要压缩的图片")
    #return 1
    if len(img_list1) == 0:
        print(mydir + "没有寻找到需要压缩的图片（不包括GIF）")
        #return 1
    else:
        with ThreadPoolExecutor(max_workers=works) as t:
            obj_list = []
            for list2 in img_list1:
                obj = t.submit(pilys2, infile=list2, code=code)
                obj_list.append(obj)
            for future in as_completed(obj_list):
                data += future.result()
    if len(gif_img_list) !=0:
        print(mydir + "含有GIF图片")
        with ThreadPoolExecutor(max_workers=works) as t:
            obj_list = []
            for list2 in gif_img_list:
                obj = t.submit(gif_img_ps, infile=list2, code=code)
                obj_list.append(obj)
            for future in as_completed(obj_list):
                data += future.result()
    if data == 0 and flag != 'no':
        write_sc(mydir, flag=flag)
    if data == 0 and os.path.exists(mydir + '\\' + after_ys):  # 删除已处理文件目录
        pass
        os.remove(mydir + '\\' + after_ys)
    if data == 0 and os.path.exists(mydir + '\\' + 'erro'):  # 删除错误文件目录
        os.remove(mydir + '\\' + 'erro')
    return data


# 搜索文件名并返回路径
def find_file(file_name, search_path):
    for root, dirs, files in os.walk(search_path):
        if file_name in files:
            return os.path.join(root, file_name)
    return ""

def winapi_path(dos_path, encoding=None):
    if (not isinstance(dos_path, str) and encoding is not None):
        dos_path = dos_path.decode(encoding)
    path = os.path.abspath(dos_path)
    if path.startswith(u"\\\\"):
        return dos_path
    return u"\\\\?\\" + path

if __name__ == '__main__':
    pass

